import torch
from transformers import BertTokenizer
from constant import *


def get_word_token_range(sentence_text: str, tokenizer: BertTokenizer) -> tuple:
    token_range, token2id, token_start = [], [], 0
    for i, word_text in enumerate(sentence_text.strip().split()):
        token_end = token_start + len(tokenizer.encode(word_text, add_special_tokens=False))
        for j in range(token_start, token_end):
            token2id.append(i)
            pass
        token_range.append([token_start, token_end - 1])
        token_start = token_end
        pass
    return token_range, token2id


class SentenceInstance:  # 某个句子的所有的原始信息都存储在这里
    def __init__(self, sentence_text: str, origin_triplets: list, tokenizer: BertTokenizer):
        self.sentence_text = sentence_text
        self.origin_triplets = origin_triplets
        self.sentence_token = tokenizer.encode(self.sentence_text)[1:-1]  #
        self.sentence_text_split_list = sentence_text.strip().split()

        self.word_token_range, self.token2id = get_word_token_range(sentence_text, tokenizer)
        # 关于这个部分的代码，具体可以看 test03.py 中的测试
        # self.word_token_range 形如 [[0, 0], [1, 1], [1, 2], [3, 4]]

        self.triplets = []
        self.generate_triplets()  # 生成上边的这个 triplets
        pass

    def generate_triplets(self):
        for aspect_span, opinion_span_list, sentiment_id in self.origin_triplets:
            triplet_dict = {"aspect_span": [], "opinion_span_list": [], 'sentiment': 0}
            triplet_dict["aspect_span"].append(self.word_token_range[aspect_span[0]][0])
            triplet_dict["aspect_span"].append(self.word_token_range[aspect_span[-1]][1])
            # 接下来是把 text 中的 aspect 和 opinion 的位置转化成在 token 中的位置
            for opinion_span in opinion_span_list:
                new_opinion_span = [self.word_token_range[opinion_span[0]][0], self.word_token_range[opinion_span[-1]][1]]
                triplet_dict["opinion_span_list"].append(new_opinion_span)
                pass
            triplet_dict['sentiment'] = sentiment_id
            self.triplets.append(triplet_dict)
            pass
        pass

    pass


class TrainingInstance:  # 用来放入模型进行训练的实例，包含着 plus_sentence 和一个 mask
    def __init__(self, sentence_instance: SentenceInstance, plus_instance: SentenceInstance, mask_str: str):
        self.sentence_instance = sentence_instance
        self.plus_instance = plus_instance
        self.mask_str = mask_str  # 就是 "0010" 类似于这样的字符串
        self.target_triplet = None  # 是一个 dict 类型 {"aspect_span": [], "opinion_span_list": [[], []], 'sentiment': 0}
        self.right_target_triplet = None  # 是一个 dict 类型 {"aspect_span": [], "opinion_span_list": [[], []], 'sentiment': 0}

        # 找到 target_triplet 和 right_target_triplet
        for char_i, char_here in enumerate(self.mask_str):
            if char_here == "1":
                self.right_target_triplet = self.sentence_instance.triplets[char_i]
                if self.target_triplet is None:
                    self.target_triplet = self.sentence_instance.triplets[char_i]
                    pass
                pass
            pass
        pass

    def update_attention_mask(self, mask_tensor: torch.Tensor, opinion_mask_str: str = None, random_mask_word: bool = False):
        # 这里的输入 tensor 是一个一维 tensor，不需要 squeeze
        sen_pre_len = len(self.plus_instance.sentence_token) + 2
        # 主要就是屏蔽掉 aspect
        for i, char_here in enumerate(self.mask_str):
            if char_here == "0":
                a_start, a_end = self.sentence_instance.triplets[i]["aspect_span"]
                mask_tensor[sen_pre_len + a_start: sen_pre_len + a_end + 1] = 0
                pass
            pass
        # 按照 opinion_mask_str 屏蔽所有的 opinion
        # 是否随机屏蔽一个单词
        pass

    def get_extended_label(self, seq_len: int = 100) -> dict:  # seq_len 是总共的长度
        sen_pre_len = len(self.plus_instance.sentence_token) + 2
        # 定义 a_s 和 a_e 和 s 和 is_on
        a_s_i = torch.tensor([sentiment2id["Invalid"]])  # -1
        a_e_i = torch.tensor([sentiment2id["Invalid"]])  # -1
        a_s_i2 = torch.tensor([sentiment2id["Invalid"]])  # -1
        a_e_i2 = torch.tensor([sentiment2id["Invalid"]])  # -1
        sentiment_id = torch.tensor([sentiment2id["Invalid"]])  # -1
        sentiment_id2 = torch.tensor([sentiment2id["Invalid"]])  # -1
        is_on = torch.tensor([0])
        # 定义 aspect_tensor
        a_s_tensor = torch.tensor([0] * (seq_len - sen_pre_len))
        a_e_tensor = torch.tensor([0] * (seq_len - sen_pre_len))
        a_tensor = torch.tensor([0] * (seq_len - sen_pre_len))
        a_s_tensor2 = torch.tensor([0] * (seq_len - sen_pre_len))
        a_e_tensor2 = torch.tensor([0] * (seq_len - sen_pre_len))
        a_tensor2 = torch.tensor([0] * (seq_len - sen_pre_len))
        # 定义 opinion_tensor
        o_s_tensor = torch.tensor([0] * (seq_len - sen_pre_len))
        o_e_tensor = torch.tensor([0] * (seq_len - sen_pre_len))
        o_tensor = torch.tensor([0] * (seq_len - sen_pre_len))
        o_strength = torch.tensor([0])
        o_s_tensor2 = torch.tensor([0] * (seq_len - sen_pre_len))
        o_e_tensor2 = torch.tensor([0] * (seq_len - sen_pre_len))
        o_tensor2 = torch.tensor([0] * (seq_len - sen_pre_len))
        o_strength2 = torch.tensor([0])
        # 正常情况：需要更新这些变量
        if self.target_triplet is not None:
            a_start, a_end = self.target_triplet["aspect_span"]
            a_start2, a_end2 = self.right_target_triplet["aspect_span"]
            # 更新 aspect 的 tensor
            a_s_tensor[a_start], a_e_tensor[a_end] = 1, 1
            a_s_tensor2[a_start2], a_e_tensor2[a_end2] = 1, 1
            a_tensor[a_start: a_end+1] = 1
            a_tensor2[a_start2: a_end2+1] = 1
            # 更新 opinion 的 tensor
            for o_start, o_end in self.target_triplet["opinion_span_list"]:
                o_s_tensor[o_start], o_e_tensor[o_end] = 1, 1
                o_tensor[o_start: o_end + 1] = 1
                pass
            o_strength = torch.tensor([len(self.target_triplet["opinion_span_list"])])
            for o_start, o_end in self.right_target_triplet["opinion_span_list"]:
                o_s_tensor2[o_start], o_e_tensor2[o_end] = 1, 1
                o_tensor2[o_start: o_end + 1] = 1
                pass
            o_strength2 = torch.tensor([len(self.target_triplet["opinion_span_list"])])
            # 更新其他内容
            a_s_i, a_e_i = torch.tensor([a_start]), torch.tensor([a_end])
            a_s_i2, a_e_i2 = torch.tensor([a_start2]), torch.tensor([a_end2])
            sentiment_id = torch.tensor([self.target_triplet["sentiment"]])
            sentiment_id2 = torch.tensor([self.right_target_triplet["sentiment"]])
            is_on = torch.tensor([1])
            pass
        # 返回最终的结果
        return {
            'a_s': a_s_i,
            'a_e': a_e_i,
            'a_s2': a_s_i2,
            'a_e2': a_e_i2,
            'a_s_': a_s_tensor.unsqueeze(0),
            'a_e_': a_e_tensor.unsqueeze(0),
            'a_t': a_tensor,
            'a_s_2': a_s_tensor2.unsqueeze(0),
            'a_e_2': a_e_tensor2.unsqueeze(0),
            'a_t2': a_tensor2,
            'is_on': is_on,
            'o_s': o_s_tensor.unsqueeze(0),
            'o_e': o_e_tensor.unsqueeze(0),
            'o_tensor': o_tensor,
            'o_strength': o_strength,
            'o_s2': o_s_tensor2.unsqueeze(0),
            'o_e2': o_e_tensor2.unsqueeze(0),
            'o_tensor2': o_tensor2,
            'o_strength2': o_strength2,
            's': sentiment_id,
            's2': sentiment_id2
        }

    pass


class TestInstance:
    def __init__(self, sentence_instance: SentenceInstance, plus_instance: SentenceInstance):
        self.meta_instance = sentence_instance
        self.plus_instance = plus_instance
        pass

    def get_token_as_batch(self, tokenizer: BertTokenizer):
        return tokenizer(
            [self.plus_instance.sentence_text],
            [self.meta_instance.sentence_text],
            padding=True,
            return_tensors='pt'
        )

    pass